/*
Copyright (c) 2008, Mike Gashler

Permission is hereby granted, free of charge, to any person
obtaining a copy of this software and associated documentation
files (the "Software"), to deal in the Software without
restriction, including without limitation the rights to use,
copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the
Software is furnished to do so, subject to the following
conditions:

The above copyright notice and this permission notice shall be
included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
OTHER DEALINGS IN THE SOFTWARE.
*/

#ifndef __GSOCKET_H__
#define __GSOCKET_H__

#include <string.h>
#include "app.h"

#ifdef WIN32
#ifndef _MT
#error ERROR: _MT not defined.  GSocket requires using the multithreaded libraries
#define _MT
#endif // !_MT
#endif // WIN32


// Other Includes
#ifdef WIN32
#	ifndef IPPROTO_IP
#		include <winsock2.h>
#	endif // IPPROTO_IP
#else // WIN32
#	include <sys/types.h>
#	include <sys/socket.h>
#	include <netinet/in.h>
#endif // else WIN32


#ifndef WIN32
typedef struct sockaddr SOCKADDR;
typedef struct sockaddr_in SOCKADDR_IN;
typedef struct hostent HOSTENT;
#endif // WIN32

class GPointerQueue;
class GSpinLock;
class GKeyPair;
class GSecureSocketServer;
class GSecureSocketClient;
class GPointerArray;
class GKeyPair;

class GDynamicArray
{
protected:
	int m_nCellSize;
	int m_nCellCount;
	int m_nAllocCount;
	int m_nGrowBy;
	void* m_pData;

public:
	// nCellSize = size of a single element in bytes.
	// nGrowBy = number of elements to initially allocate space
	//           for when the first element is added to the array
	GDynamicArray(int nCellSize, int nGrowBy);
	virtual ~GDynamicArray();

#ifndef NO_TEST_CODE
	static void Test();
#endif // NO_TEST_CODE
	// Returns a reference to element n
	inline void* _GetCellRef(int n)
	{
		if(n < 0 || n >= m_nCellCount)
			ThrowError("index out of range");
		return((char*)m_pData + (m_nCellSize * n));
	}

	// Adds an element (referenced by pData) to the end of the array.
	void _AddCellByRef(const void* pData);

	// Replaces element nCell with the element reference by pData
	void _SetCellByRef(int nCell, void* pData);

	// Deletes an element and moves all subsequent elements over
	// to fill the gap.  (If order is not important in your
	// array, it's a good idea to swap the cell you want to
	// delete with the last element before you delete
	// it so that you won't have to pay the performance
	// penalty for shifting the whole array over.)
	void DeleteCell(int n);

	// Inserts an element (referenced by pData) into the array at location n.
	// (It will shift all subsequent elements over to make room.)
	void _InsertCellByRef(int n, void* pData);

	// Returns the number of elements in the array
	int GetSize() { return m_nCellCount; }

	// Resizes the array (throwing out elements >= n)
	void SetSize(int n);

	// Returns the number of bytes occupied by each cell
	int GetCellSize() { return m_nCellSize; }

	// Resizes the array to zero elements
	void Clear() { SetSize(0); }

	void SetAllocSize(int n);
	int GetAllocSize() { return m_nAllocCount; }
	int GetGrowBy() { return m_nGrowBy; }

	// Make the allocated buffer bigger
	void Grow();

	void* GetRawData() { return m_pData; }
};

class GPointerArray : public GDynamicArray
{
public:
	GPointerArray(int nGrowBy) : GDynamicArray(sizeof(void*), nGrowBy) { }
	virtual ~GPointerArray() { }

	inline void* GetPointer(int nIndex) { return *(void**)_GetCellRef(nIndex); }
	inline void AddPointer(const void* p) { _AddCellByRef(&p); }
	void InsertPointer(int nPos, void* p) { _InsertCellByRef(nPos, &p); }
	inline void SetPointer(int nCell, void* p) { _SetCellByRef(nCell, &p); }

	// Sorts the array with merge sort
	//void Sort(PointerComparer pCompareFunc, void* pThis);
};

class GIntArray : public GDynamicArray
{
public:
	GIntArray(int nGrowBy) : GDynamicArray(sizeof(int), nGrowBy) { }
	virtual ~GIntArray() { }

	int GetInt(int nIndex) { return *(int*)_GetCellRef(nIndex); }
	void AddInt(int n) { _AddCellByRef(&n); }
	void InsertInt(int nPos, int nVal) { _InsertCellByRef(nPos, &nVal); }
	void SetInt(int nCell, int n) { _SetCellByRef(nCell, &n); }
};


class GSocketArray : public GDynamicArray
{
public:
	GSocketArray(int nGrowBy) : GDynamicArray(sizeof(SOCKET), nGrowBy) { }
	virtual ~GSocketArray() { }

	SOCKET GetSocket(int nIndex) { return *(SOCKET*)_GetCellRef(nIndex); }
	void AddSocket(SOCKET s) { _AddCellByRef(&s); }
	void SetSocket(int nCell, SOCKET s) { _SetCellByRef(nCell, &s); }
};




class GSocketClientBase
{
protected:
	bool m_bKeepListening;
	SOCKET m_s;
	HANDLE m_hListenThread;
	bool m_bUDP;

public:
	GSocketClientBase(bool bUDP);
	virtual ~GSocketClientBase();

	static bool IsThisAnIPAddress(const char* szHost);

	// This returns the SOCKET
	SOCKET GetSocketHandle();

	// nTimeout is specified in seconds
	bool Connect(const char* szHost, unsigned short nPort, int nTimeout = 10);
	void Disconnect();

	bool Send(const unsigned char *pBuf, int nLen);

	// This method is abstract because you need to implement something here
	virtual bool Receive(unsigned char *pBuf, int nLen) = 0; // Override me


	u_short GetMyPort();
	struct in_addr GetMyIPAddr();
	bool IsConnected();
	char* GetMyIPAddr(char* szBuff, int nBuffSize);
	char* GetMyName(char* szBuff, int nBuffSize);
	u_short GetTheirPort();
	struct in_addr GetTheirIPAddr();
	char* GetTheirIPAddr(char* szBuff, int nBuffSize);
	char* GetTheirName(char* szBuff, int nBuffSize);

	void Listen(); // Don't call this method directly

	// This parses a URL into its parts
	static void ParseURL(const char* szUrl, int* pnHostIndex, int* pnPortIndex, int* pnPathIndex, int* pnParamsIndex);

	// Parses the parameter portion of a URL
	static int ParseUrlParams(const char* szParams, int nMaxParams, char** pNames, int* pNameLengths, char** pValues, int* pValueLengths);

protected:
	bool GoIntoHostMode(unsigned short nListenPort, int nMaxConnections);
	int GetFirstAvailableSocketNumber();
	void JoinListenThread();
	void JoinListenThread(int nConnectionNumber);

	// This method is empty. It's just here so you can override it.
	// This is called when the connection is gracefully closed. (There is no built-in support
	// for detecting ungraceful disconnects. This is a feature, not a bug, because it makes
	// it robust to sporadic hardware. I recommend implementing a system where the
	// server requires the client to send periodic heartbeat packets and you call Disconnect()
	// if the responses don't come regularly.)
	virtual void OnCloseConnection();
};








class GSocketServerBase
{
protected:
	SOCKET m_socketConnectionListener;
	GSocketArray* m_pConnections;
	GSpinLock* m_pConnectionsLock;
	fd_set m_socketSet; // structure used by select()
	HANDLE m_hWorkerThread;
	bool m_bKeepWorking;
	SOCKADDR_IN m_sHostAddrIn;
	bool m_bUDP;
	int m_nMaxConnections;
	char* m_szReceiveBuffer;

public:
	GSocketServerBase(bool bUDP, int nPort, int nMaxConnections);
	virtual ~GSocketServerBase();

	// This returns the SOCKET
	SOCKET GetSocketHandle(int nConnectionNumber);

	void Disconnect(int nConnectionNumber);

	bool Send(const unsigned char *pBuf, int nLen, int nConnectionNumber);

	// This method is abstract because you need to implement something here
	virtual bool Receive(unsigned char *pBuf, int nLen, int nConnectionNumber) = 0; // Override me

	bool IsConnected(int nConnectionNumber);

	// Stops listening for new connections
	void DontAcceptAnyMoreConnections();

	// If nConnectionNumber is less than 0, returns the address of the server
	char* GetIPAddr(char* szBuff, int nBuffSize, int nConnectionNumber);

	// If nConnectionNumber is less than 0, returns the address of the server
	struct in_addr GetIPAddr(int nConnectionNumber);

	// If nConnectionNumber is less than 0, returns the port of the server
	u_short GetPort(int nConnectionNumber);

	// If nConnectionNumber is less than 0, returns the name of the server
	char* GetName(char* szBuff, int nBuffSize, int nConnectionNumber);

	void ServerWorker(); // Don't call this method directly

protected:
	void Init(bool bUDP, int nPort, int nMaxConnections);
	int GetFirstAvailableConnectionNumber();
	void JoinWorkerThread();

	// This method is empty. It's just here so you can override it.
	// This is called when the connection is gracefully closed. (There is no built-in support
	// for detecting ungraceful disconnects. This is a feature, not a bug, because it makes
	// it robust to sporadic hardware. To detect ungraceful disconnects, I recommend requiring
	// the client to send periodic heartbeat packets and calling Disconnect() if they stop coming.)
	virtual void OnCloseConnection(int nConnection);

	// This method is empty. It's just here so you can override it.
	// WARNING: the connection isn't fully open at the time this method is called,
	//          so don't send anything back to the client inside this callback
	virtual void OnAcceptConnection(int nConnection);

	SOCKET RefreshSocketSet();
	int HandleNewConnection();
	void ReduceConnectionList();
};

// --------------------------------------------------------------------------

struct GEZSocketPacketHeader
{
	char tag[4];
	int nPayloadSize;
};


class GSocketMessage
{
protected:
	unsigned char* m_pMessage;
	int m_nMessageSize;
	int m_nConnection;

public:
	GSocketMessage(unsigned char* pMessage, int nMessageSize, int nConnection)
	{
		m_pMessage = new unsigned char[nMessageSize];
		memcpy(m_pMessage, pMessage, nMessageSize);
		m_nMessageSize = nMessageSize;
		m_nConnection = nConnection;
	}

	virtual ~GSocketMessage()
	{
		delete(m_pMessage);
	}

	const unsigned char* GetTheMessage() { return m_pMessage; }
	int GetMessageSize() { return m_nMessageSize; }
	int GetConnection() { return m_nConnection; }

	// you must delete the buffer this returns
	unsigned char* TakeBuffer()
	{
		unsigned char* pMessage = m_pMessage;
		m_pMessage = NULL;
		return pMessage;
	}
};


// --------------------------------------------------------------------------


// This class is designed to make network communication easy
class GSocketServer : public GSocketServerBase
{
protected:
	GPointerArray* m_pBuffers;
	int m_nMaxPacketSize;
	GPointerQueue* m_pMessageQueue;
	GSpinLock* m_pMessageQueueLock;

	virtual bool Receive(unsigned char *pBuf, int len, int nConnectionNumber);
	void QueueMessage(unsigned char* pBuf, int nLen, int nConnectionNumber);

	GSocketServer(bool bUDP, int nMaxPacketSize, int nPort, int nMaxConnections);

public:
	// Host a TCP socket
	static GSocketServer* HostTCPSocket(int nPort);

	// Host a UDP socket
	static GSocketServer* HostUDPSocket(int nPort);

	// Host a Gash socket.  (Adds a layer on top of TCP to guarantee same-size packet delivery)
	static GSocketServer* HostGashSocket(int nPort, int nMaxPacketSize);

	virtual ~GSocketServer();

	// Send some data
	bool Send(const void* pBuf, int nLen, int nConnectionNumber);

	// Returns the number of messages waiting to be received
	int GetMessageCount();

	// Receive the next message. (You are responsible to delete the buffer this returns)
	unsigned char* GetNextMessage(int* pnSize, int* pnOutConnectionNumber);
};

// --------------------------------------------------------------------------

// This class is designed to make network communication easy
class GSocketClient : public GSocketClientBase
{
protected:
	unsigned char* m_pBuffer;
	int m_nBufferPos;
	int m_nMaxPacketSize;
	GPointerQueue* m_pMessageQueue;
	GSpinLock* m_pMessageQueueLock;

	virtual bool Receive(unsigned char *pBuf, int len);
	void QueueMessage(unsigned char* pBuf, int nLen);

	GSocketClient(bool bUDP, int nMaxPacketSize);

public:
	// Connect to a TCP socket
	static GSocketClient* ConnectToTCPSocket(const char* szAddress, int nPort);

	// Connect to a UDP socket
	static GSocketClient* ConnectToUDPSocket(const char* szAddress, int nPort);

#ifndef NO_TEST_CODE
	static void Test();
#endif // !NO_TEST_CODE

	// Connect to a Gash socket.  (Adds a layer on top of TCP to guarantee same-size packet delivery)
	static GSocketClient* ConnectToGashSocket(const char* szAddress, int nPort, int nMaxPacketSize);

	virtual ~GSocketClient();

	// Send some data
	bool Send(const void* pBuf, int nLen);

	// Returns the number of messages waiting to be received
	int GetMessageCount();

	// Receive the next message. (You are responsible to delete the buffer this returns)
	unsigned char* GetNextMessage(int* pnSize);
};


#endif // __GSOCKET_H__
